import networkx as nx
import numpy as np
import random
import abc
import pdb

# 这是一个抽象基类（Abstract Base Class），定义了 gen_node_features 抽象方法，该方法在子类中需要被实现。这个类是一个特征生成器的基类，所有的特征生成器都应该继承自这个类并实现 gen_node_features 方法。
# 定义了一个名为 FeatureGen 的类，并指定了元类为 abc.ABCMeta。这表示 FeatureGen 是一个抽象基类（Abstract Base Class），它提供了一个接口以及一些方法的默认实现，但不会在这里实现具体的方法。
class FeatureGen(metaclass=abc.ABCMeta):
    """Feature Generator base class."""
    # @abc.abstractmethod：使用装饰器 @abc.abstractmethod 标记了一个抽象方法 gen_node_features()。这意味着任何继承自 FeatureGen 的子类必须实现该方法。
    @abc.abstractmethod
    # 定义了一个抽象方法 gen_node_features()，该方法接受一个参数 G，用于生成节点特征。
    def gen_node_features(self, G):
        pass

# 定义了一个名为 ConstFeatureGen 的类，它继承自 FeatureGen 类。
class ConstFeatureGen(FeatureGen):
    """Constant Feature class."""
    # 定义了 ConstFeatureGen 类的构造函数。它接受两个参数 val 和 max_degree，其中 val 是特征值，max_degree 是节点的最大度数。构造函数初始化了类的属性 val、max_degree 和 one_hots。
    def __init__(self, val, max_degree=20):
        self.val = val
        self.max_degree = max_degree
        # 创建了一个大小为 (max_degree, max_degree) 的单位矩阵，并将其赋值给 one_hots 属性。这个矩阵用于生成节点的 one-hot 特征表示。
        self.one_hots = np.eye(max_degree, dtype=np.float32)

    def gen_node_features(self, G):
        # 创建了一个空字典，用于存储生成的节点特征。
        feat_dict = {}
        # 检查 val 是否为 None。如果是 None，则表示要根据节点的度数生成特征；否则，表示所有节点都具有相同的特征值。
        if self.val is None:
            # 遍历图 G 中的节点和它们的度数。
            for k, v in G.degree():
                # 检查节点的度数是否小于 max_degree。
                if v < self.max_degree:
                    # 将节点 k 的特征设置为对应度数的 one-hot 编码。
                    feat_dict[k] = {'feat': self.one_hots[v]}
                else:
                    # 将节点 k 的特征设置为度数为 max_degree - 1 的 one-hot 编码。
                    feat_dict[k] = {'feat': self.one_hots[self.max_degree - 1]}  
        else:
            # 如果 val 不是 None，表示所有节点都具有相同的特征值。
            # 创建一个字典，其中键是节点的索引，值是具有相同特征值的特征向量。
            feat_dict = {i:{'feat': np.array(self.val, dtype=np.float32)} for i in G.nodes()}
        # 使用生成的特征字典将特征添加到图 G 的节点属性中。
        nx.set_node_attributes(G, feat_dict)


# 定义了一个名为 GaussianFeatureGen 的类，它继承自 FeatureGen 类。
class GaussianFeatureGen(FeatureGen):
    """Gaussian Feature class."""
    # 定义了 GaussianFeatureGen 类的构造函数。它接受两个参数 mu 和 sigma，分别表示高斯分布的均值向量和协方差矩阵。
    def __init__(self, mu, sigma):
        # 将输入的均值向量赋值给类的属性 mu。
        self.mu = mu
        # 检查协方差矩阵 sigma 的维度是否小于 2。
        if sigma.ndim < 2:
            # 如果协方差矩阵是一维的，将其转换为对角矩阵。这是为了处理输入的协方差向量的情况。
            self.sigma = np.diag(sigma)
        else:
            # 如果协方差矩阵的维度大于等于 2。
            # 直接将输入的协方差矩阵赋值给类的属性 sigma。
            self.sigma = sigma

    def gen_node_features(self, G):
        # 使用 numpy.random.multivariate_normal() 函数生成服从给定均值向量和协方差矩阵的多元正态分布的特征向量。G.number_of_nodes() 指定了要生成的特征向量的数量，即图 G 中的节点数。
        feat = np.random.multivariate_normal(self.mu, self.sigma, G.number_of_nodes())
        # 创建一个字典 feat_dict，其中键是节点的索引，值是对应节点的特征向量。
        feat_dict = {
                i: {"feat": feat[i]} for i in range(feat.shape[0])
            }
        # 使用生成的特征字典将特征添加到图 G 的节点属性中。
        nx.set_node_attributes(G, feat_dict)

# 定义了一个名为 GridFeatureGen 的类，它继承自 FeatureGen 类。
class GridFeatureGen(FeatureGen):
    """Grid Feature class."""
    # 定义了 GridFeatureGen 类的构造函数。它接受三个参数 mu、sigma 和 com_choices，分别表示特征的均值、方差以及可能的社区标签列表。
    def __init__(self, mu, sigma, com_choices):
        # 将输入的均值、方差和社区标签列表分别赋值给类的属性 mu、sigma 和 com_choices。
        self.mu = mu                    # Mean
        self.sigma = sigma              # Variance
        self.com_choices = com_choices  # List of possible community labels

    def gen_node_features(self, G):
        # Generate community assignment
        # 根据节点的度数判断其所属的社区标签，并将结果存储在 community_dict 字典中。如果节点的度数小于 4，则属于 com_choices 列表的第一个社区标签，否则属于列表的第二个社区标签。
        community_dict = {
            n: self.com_choices[0] if G.degree(n) < 4 else self.com_choices[1]
            for n in G.nodes()
        }

        # Generate random variable
        # 使用 numpy.random.normal() 函数生成服从给定均值和方差的正态分布的随机变量 s，其中 G.number_of_nodes() 指定了要生成的随机变量的数量，即图 G 中的节点数。
        s = np.random.normal(self.mu, self.sigma, G.number_of_nodes())

        # Generate features
        # 创建一个字典 feat_dict，其中键是节点的索引，值是包含社区标签和相应随机变量值的特征向量。
        feat_dict = {
            n: {"feat": np.asarray([community_dict[n], s[i]])}
            for i, n in enumerate(G.nodes())
        }
        # 使用生成的特征字典将特征添加到图 G 的节点属性中，并返回社区分配字典 community_dict。
        nx.set_node_attributes(G, feat_dict)
        return community_dict

